#include "block_hhmm.hpp"
#include "hypercube/hypercube_cryto_agent.hpp"
#include "matrix_util.hpp"
#include "mpi//get_matrix_blk.hpp"
#include <gtest/gtest.h>
#include <iostream>

static params<helib::BGV> my_params;
bool check = false;
bool no_encoding_reuse = false;
unsigned long input_bit_len;

#define STDBLK_CGEMM_TEST(DIM, BLK)                                            \
  TEST(TEST_STDBLK, MNK##DIM##_BLK##BLK) {                                     \
    unsigned long blk_size = BLK;                                              \
    unsigned long m = DIM;                                                     \
    unsigned long n = DIM;                                                     \
    unsigned long k = DIM;                                                     \
                                                                               \
    /* Initialize data */                                                      \
    int *A = nullptr;                                                          \
    int *B = nullptr;                                                          \
    int *C = nullptr;                                                          \
                                                                               \
    FIX_RAND_SEED(true);                                                       \
    const std::size_t size_A = m * k;                                          \
    const std::size_t size_B = k * n;                                          \
    ALLOC_INT_ARRAY_WITHOUT_INIT(A, size_A);                                   \
    ALLOC_INT_ARRAY_WITHOUT_INIT(B, size_B);                                   \
    RAND_ARRAY(A, size_A, input_bit_len);                                      \
    RAND_ARRAY(B, size_B, input_bit_len);                                      \
    if (check) {                                                               \
      const std::size_t size_C = m * n;                                        \
      ALLOC_INT_ARRAY_WITHOUT_INIT(C, size_C);                                 \
      cblas_igemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0, A,  \
                  k, B, n, 0.0, C, n);                                         \
    }                                                                          \
                                                                               \
    /* Initialize cryptograph system and HMM engine */                         \
    const auto eh_cube =                                                       \
        hypercube_hmme::get_hypercube(blk_size, blk_size, blk_size);           \
    hypercube_hmmcc hmmcc(my_params, eh_cube, true);                           \
    hmmcc.report_context();                                                    \
                                                                               \
    hmm_status<shmm_engine> status(METHOD::BSGS, omp_get_max_threads());       \
    hypercube_hmme hmme(status);                                               \
                                                                               \
    const auto eh_mnk =                                                        \
        hypercube_hmme::get_expected_mnk(blk_size, blk_size, blk_size);        \
    hmme.register_engine(hmmcc.context()->getEA(), eh_mnk);                    \
                                                                               \
    /* Matrix Partition */                                                     \
    unsigned long m_nblk = m / blk_size;                                       \
    unsigned long n_nblk = n / blk_size;                                       \
    unsigned long k_nblk = k / blk_size;                                       \
    /* For Matrix A */                                                         \
    mat_ctxt ctxt_A;                                                           \
    ctxt_A.resize(m_nblk);                                                     \
    for (std::size_t i = 0; i < m_nblk; i++) {                                 \
      for (std::size_t j = 0; j < k_nblk; j++) {                               \
        auto mat = get_blk(A, m, k, i, j, blk_size, blk_size);                 \
        ctxt_A[i].emplace_back(hmmcc.encrypt(mat));                            \
      }                                                                        \
    }                                                                          \
    /* For Matrix B */                                                         \
    mat_ctxt ctxt_B;                                                           \
    ctxt_B.resize(k_nblk);                                                     \
    for (std::size_t i = 0; i < k_nblk; i++) {                                 \
      for (std::size_t j = 0; j < n_nblk; j++) {                               \
        auto mat = get_blk(B, k, n, i, j, blk_size, blk_size);                 \
        ctxt_B[i].emplace_back(hmmcc.encrypt(mat));                            \
      }                                                                        \
    }                                                                          \
    /* Encoding Reuse Close */                                                 \
    if (no_encoding_reuse) {                                                   \
      hmme.status()->reset(false);                                             \
      hmme.status()->set_encoding_reuse(false);                                \
    }                                                                          \
                                                                               \
    metrics metric;                                                            \
    mat_ctxt ctxt_C;                                                           \
    auto type = shmm_engine::parse_matrix_type(eh_mnk);                        \
    blk_matmul(ctxt_C, ctxt_A, ctxt_B, hmme, type, metric);                    \
    if (no_encoding_reuse) {                                                   \
      std::cout << "* HMM encoding time : " << hmme.get_encoding_time()        \
                << " s\n";                                                     \
    }                                                                          \
    metric.report();                                                           \
                                                                               \
    if (!check) {                                                              \
      delete[] A;                                                              \
      delete[] B;                                                              \
      GTEST_SKIP() << "Skip checking answers";                                 \
    }                                                                          \
    /* Check answers */                                                        \
    for (std::size_t i = 0; i < m_nblk; i++) {                                 \
      for (std::size_t j = 0; j < n_nblk; j++) {                               \
        NTL::mat_ZZ res;                                                       \
        hmmcc.decrypt(res, ctxt_C[i][j]);                                      \
        for (std::size_t ii = 0; ii < blk_size; ii++) {                        \
          for (std::size_t jj = 0; jj < blk_size; jj++) {                      \
            const auto expect =                                                \
                C[(i * blk_size + ii) * n + (j * blk_size + jj)];              \
            const auto actual = res[ii][jj];                                   \
            ASSERT_EQ(expect, actual)                                          \
                << "diff at (" << ii << ", " << jj << ") in BLK (" << i        \
                << ", " << j << ")";                                           \
          }                                                                    \
        }                                                                      \
      }                                                                        \
    }                                                                          \
    delete[] A;                                                                \
    delete[] B;                                                                \
    delete[] C;                                                                \
  }

STDBLK_CGEMM_TEST(64, 32)
STDBLK_CGEMM_TEST(256, 128)
STDBLK_CGEMM_TEST(512, 128)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);

  command_line_parser parser(my_params);
  parser.add_option("-chk", "enable_check_answer", &check);
  parser.add_option("-ner", "no_encoding_reuse", &no_encoding_reuse);
  parser.add_option("-ib", "7", "input_bit_length", &input_bit_len);
  int ret = parser.parse(argc, argv);
  EXPECT_EQ(ret, 0);

  // Run tests
  ret = RUN_ALL_TESTS();

  return ret;
}